import com.google.common.annotations.VisibleForTesting;
import java.util.concurrent.locks.ReentrantLock;

/** Token bucket for limiting hedged requests. Thread-safe. */
final class HedgeTokenBucket {
  private final ReentrantLock lock = new ReentrantLock();
  private final int maxTokens;
  private final float refillAmount;
  private float tokens;

  HedgeTokenBucket(HedgeSettings settings) {
    this.maxTokens = settings.getMaxTokens();
    this.refillAmount = settings.getRefill();
    this.tokens = maxTokens;
  }

  @VisibleForTesting
  HedgeTokenBucket(int maxTokens, float refillAmount) {
    this.maxTokens = maxTokens;
    this.refillAmount = refillAmount;
    this.tokens = maxTokens;
  }

  /**
   * Attempts to acquire a token for a hedged request.
   *
   * @return {@code true} if a token was acquired, {@code false} otherwise.
   */
  boolean tryAcquire() {
    lock.lock();
    try {
      if (tokens < 1.0f) {
        return false;
      }
      tokens -= 1.0f;
      return true;
    } finally {
      lock.unlock();
    }
  }

  /** Refills the bucket by the configured refill amount, capped at max tokens. */
  void refill() {
    lock.lock();
    try {
      tokens = Math.min((float) maxTokens, tokens + refillAmount);
    } finally {
      lock.unlock();
    }
  }

  @VisibleForTesting
  float getTokens() {
    lock.lock();
    try {
      return tokens;
    } finally {
      lock.unlock();
    }
  }